package com.itcam.chat.service.impl;

import cn.hutool.core.util.IdUtil;
import com.itcam.ai.utils.MongoUtil;
import com.itcam.chat.domain.pojo.Message;
import com.itcam.chat.domain.vo.MessageVo;
import com.itcam.chat.domain.vo.QueryVo;
import com.itcam.chat.service.OllamaService;
import com.itcam.common.utils.uuid.IdUtils;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.vectorstore.RedisVectorStore;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.Date;
import java.util.List;

/**
 * @author : Cammy.Wu
 *
 */

@Service
public class OllamaServiceImpl implements OllamaService {

    @Autowired
    private MongoTemplate mongoTemplate;

    @Autowired
    private OllamaChatModel ollamaChatModel;

    @Autowired
    private RedisVectorStore redisVectorStore;

    /**
     * 文本问答
     *
     *  1.获取用户提问信息并保存到mongodb中
     *  2.先查看本地知识库 将本地知识库内容加入本地系统提示中
     *  3.在将用户提问信息本身+本地知识库一起提交到Ollama qen2:7b模型进行回复
     *
     * @param queryVo
     * @return
     */
    @Override
    public Flux<String> chatStream(QueryVo queryVo) {
        // 1.将问题记录到mongodb中
        Long chatId = queryVo.getChatId();
        String msg = queryVo.getMsg();
        Long projectId = queryVo.getProjectId();
        if (chatId != null) {
            Message message = new Message();
            message.setId(IdUtil.getSnowflake().nextId());
            message.setChatId(chatId);
            message.setType(0);
            message.setContent(msg);
            message.setCreateTime(new Date());
            // 将message对象插入到MongoDB数据库的集合中要插入到哪个集合中由queryVo对象中的chatId决定
            mongoTemplate.insert(message, MongoUtil.getMessageCollection(chatId));
        }

        // 2.查询本地知识库
        // 创建一个查询请求，并设置过滤条件来筛选出符合特定projectId的记录
        List<Document> documents = redisVectorStore.similaritySearch(SearchRequest.query(msg)
                .withFilterExpression(new FilterExpressionBuilder().eq("projectId", projectId).build())
                .withSimilarityThreshold(0.7d)
                .withTopK(10));
        // 将本地知识库的内容作为系统内容提前放入
        List listMsg = new ArrayList<>();
        documents.forEach(item -> {
            String content = item.getContent();
            SystemMessage systemMessage = new SystemMessage(content);
            listMsg.add(systemMessage);
        });
        // 加入当前用户的提问
        listMsg.add(new SystemMessage(msg));

        // 3.提交到大模型获取最终结果
        Prompt prompt = new Prompt(listMsg);
        Flux<ChatResponse> responseFlux = ollamaChatModel.stream(prompt);

        // 使用map操作符对responseFlux流中的每个ChatResponse对象进行转换
        // map操作符将每个元素映射到新的值，这里是处理响应的内容
        return responseFlux.map(response -> {
            if (response.getResult().getOutput().getContent() != null) {
                return response.getResult().getOutput().getContent();
            }
            return "";
        });
    }

    /**
     * 保存AI回答的结果
     * @param messageVo
     */
    @Override
    public void saveMsg(MessageVo messageVo) {
        Message message = new Message();
        BeanUtils.copyProperties(messageVo, message);
        message.setType(1);
        message.setId(IdUtil.getSnowflake().nextId());
        message.setCreateTime(new Date());
        // 保存mongodb中
        mongoTemplate.insert(message, MongoUtil.getMessageCollection(messageVo.getChatId()));
    }
}
